/* Part of the KnowARC Janitor Use-case processor for taverna
 *  written 2007-2010 by Hajo Nils Krabbenhoeft and Steffen Moeller
 *  University of Luebeck, Institute for Neuro- and Bioinformatics
 *  University of Luebeck, Institute for Dermatolgy
 *
 *  This package is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation; either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This package is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this package; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301 USA
 */

package de.uni_luebeck.inb.knowarc.usecases.invocation.ssh;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.Vector;

import org.globus.ftp.exception.NotImplementedException;

import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.ChannelSftp;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.SftpException;
import com.jcraft.jsch.ChannelSftp.LsEntry;

import de.uni_luebeck.inb.knowarc.usecases.ScriptOutput;
import de.uni_luebeck.inb.knowarc.usecases.UseCaseDescription;
import de.uni_luebeck.inb.knowarc.usecases.invocation.AskUserForPw;
import de.uni_luebeck.inb.knowarc.usecases.invocation.UseCaseInvocation;

/**
 * The job is executed by connecting to a worker pc using ssh, i.e. not via the
 * grid.
 * 
 * @author Hajo Krabbenhoeft
 */
public class SshUseCaseInvocation extends UseCaseInvocation {
	
	private static JSch jsch = new JSch();
	
	private static Map<String, Session> sessionMap = Collections.synchronizedMap(new HashMap<String, Session> ());
	private static Map<Session, ChannelSftp> sftpMap = Collections.synchronizedMap(new HashMap<Session, ChannelSftp> ());

	private final Session sshSession;
	private final ChannelSftp sftp;
	private String tmpname;
	private final SshNode workerNode;
	private final AskUserForPw askUserForPw;
	
	private ChannelExec running;
	private final ByteArrayOutputStream stdout_buf = new ByteArrayOutputStream();
	private final ByteArrayOutputStream stderr_buf = new ByteArrayOutputStream();


	public static String test(final SshNode workerNode, final AskUserForPw askUserForPw) {
		try {
			Session sshSession = getSshSession(workerNode, askUserForPw);

			ChannelSftp sftp = (ChannelSftp) sshSession.openChannel("sftp");
			sftp.connect();
			sftp.cd(workerNode.getDirectory());
		} catch (JSchException e) {
			return e.toString();
		} catch (SftpException e) {
			return e.toString();
		}
		return null;
	}

	private static Session getSshSession(final SshNode workerNode, final AskUserForPw askUserForPw) throws JSchException {
		if (!workerNode.getDirectory().endsWith("/")) {
			workerNode.setDirectory(workerNode.getDirectory() + "/");
		}
		return getSshSession(workerNode.getHost(), workerNode.getPort(), workerNode.getDirectory(), askUserForPw);
	}
	
	private static synchronized Session getSshSession(final String host, final int port, final String directory, final AskUserForPw askUserForPw) throws JSchException {

		String sshUrl = createSshUrl(host, port, directory);

		Session s = sessionMap.get(sshUrl);
		if ((s != null) && s.isConnected()) {
			System.err.println("Reusing session");
			return s;
		}
		if (s != null) {
			System.err.println("Session was not connected");
		}
		if (s == null) {
			System.err.println("No session found for " + sshUrl);
		}

		if (askUserForPw.getKeyfile().length() > 0) {
			jsch.addIdentity(askUserForPw.getKeyfile());
		}
		Session sshSession = jsch.getSession(askUserForPw.getUsername(), host, port);
		sshSession.setUserInfo(new SshAutoLoginTrustEveryone(askUserForPw));
		sshSession.connect();
		
		askUserForPw.authenticationSucceeded();
		sessionMap.put(sshUrl, sshSession);
		if (sshSession == null) {
			System.err.println("Returning a null session");
		}
		return sshSession;
	}
	
	private static synchronized ChannelSftp getSftpChannel(Session session) throws JSchException {
		ChannelSftp result = sftpMap.get(session);
		if (result == null) {
			System.err.println("Creating new sftp channel");
			result = (ChannelSftp) session.openChannel("sftp");
			sftpMap.put(session, result);
		}
		else {
			System.err.println("Reusing sftp channel");			
		}
		if (!result.isConnected()) {
			System.err.println("Connecting");
			result.connect();
		} else {
			System.err.println("Already connected");
		}
		return result;
	}

	public static String createSshUrl(String host, int port, String directory) {
		return ("ssh://" + host + ":" + port + "/" + directory);
	}

	public SshUseCaseInvocation(UseCaseDescription desc, SshNode workerNodeA, AskUserForPw askUserForPwA)
			throws JSchException, SftpException {
		this.workerNode = workerNodeA;
		this.askUserForPw = askUserForPwA;
		usecase = desc;

		try {
			sshSession = getSshSession(workerNode, askUserForPw);
		} catch (JSchException e) {
			e.printStackTrace();
			throw e;
		}
		
		sftp = getSftpChannel(sshSession);
		synchronized(sftp) {

			sftp.cd(workerNode.getDirectory());
			Random rnd = new Random();
			while (true) {
				tmpname = "usecase" + rnd.nextLong();
				try {
					sftp.lstat(workerNode.getDirectory() + tmpname);
					continue;
				} catch (Exception e) {
					// file seems to not exist :)
				}
				sftp.mkdir(workerNode.getDirectory() + tmpname);
				sftp.cd(workerNode.getDirectory() + tmpname);
				break;
			}
		}
	}

	@Override
	public void putFile(String name, byte[] contents) {
		synchronized(sftp) {
			try {
				sftp.cd(workerNode.getDirectory() + tmpname);
			} catch (SftpException e1) {
				System.err.println("Unable to change directory" + e1);
			}
			try {
				sftp.put(new ByteArrayInputStream(contents), name);
			} catch (Exception e) {
				// TODO
			}
		}
	}

	@Override
	public void putReference(String name, String source) throws NotImplementedException {
		throw new NotImplementedException();
	}

	private void recursiveDelete(String path) throws SftpException {
		if (!path.startsWith(workerNode.getDirectory() + tmpname))
			return;
		Vector<?> entries = sftp.ls(path);
		for (Object object : entries) {
			LsEntry entry = (LsEntry) object;
			if (entry.getFilename().equals(".") || entry.getFilename().equals(".."))
				continue;
			if (entry.getAttrs().isDir())
				recursiveDelete(path + entry.getFilename() + "/");
			else
				sftp.rm(path + entry.getFilename());
		}
		sftp.rmdir(path);
	}

	@Override
	public void Cleanup() {
		try {
//			sftp.cd(workerNode.getDirectory());
//			recursiveDelete(workerNode.getDirectory() + tmpname + "/");
		} catch (Exception e) {
		    // TODO
		}
		try {
			if (running != null) {
				running.disconnect();
			}
//			sftp.disconnect();
//			sshSession.disconnect();
		} catch (Exception e) {
		}
	}

	/**
	 * Transforms an input stream towards an entity that Taverna can work with.
	 * 
	 * @param read
	 * @param binary
	 * @return
	 * @throws IOException
	 */
	private Object file2data(InputStream read, boolean binary) throws IOException {
		byte[] data = new byte[0];
		byte[] buf = new byte[1024];
		int r;
		while (-1 != (r = read.read(buf))) {
			byte[] d2 = new byte[data.length + r];
			System.arraycopy(data, 0, d2, 0, data.length);
			System.arraycopy(buf, 0, d2, data.length, r);
			data = d2;
		}
		if (binary)
			return data;
		else
			return Charset.forName("US-ASCII").decode(ByteBuffer.wrap(data)).toString();
	}

	@Override
	protected void submit_generate_job_inner() throws Exception {
		tags.put("uniqueID", "" + getSubmissionID());
		String command = usecase.getCommand();
		for (String cur : tags.keySet()) {
			command = command.replaceAll("\\Q%%" + cur + "%%\\E", tags.get(cur));
		}
		command = "cd " + workerNode.getDirectory() + tmpname + " && " + command;
		
		running = (ChannelExec) sshSession.openChannel("exec");
		running.setCommand(command);
		running.setOutputStream(stdout_buf);
		running.setErrStream(stderr_buf);
		running.connect();
	}

	@Override
	public HashMap<String, Object> submit_wait_fetch_results() throws Exception {
		while (!running.isClosed()) {
			try {
				Thread.sleep(1000);
			} catch (Exception e) {
			}
		}

		int exitcode = running.getExitStatus();
		if (exitcode != 0)
			throw new Exception("nonzero exit code");

		HashMap<String, Object> results = new HashMap<String, Object>();
		results.put("STDOUT", Charset.forName("US-ASCII").decode(ByteBuffer.wrap(stdout_buf.toByteArray())).toString());
		results.put("STDERR", Charset.forName("US-ASCII").decode(ByteBuffer.wrap(stderr_buf.toByteArray())).toString());

		synchronized(sftp) {
				sftp.cd(workerNode.getDirectory() + tmpname);
				for (Map.Entry<String, ScriptOutput> cur : usecase.getOutputs().entrySet()) {
					results.put(cur.getKey(), file2data(sftp.get(cur.getValue().getPath()), cur.getValue().isBinary()));
				}
		}

		return results;
	}

}
